#include <gtest/gtest.h>

#include "common_components/common_runtime/src/heap/allocator/treap.h"


using namespace panda;

class TreapTest : public ::testing::Test {
protected:
    void SetUp() override {
        treap_.Init(1024); // 初始化足够大的内存块
    }

    void TearDown() override {
        treap_.Fini();
    }

    Treap treap_;
};

// 测试初始化后 Treap 是否为空
TEST_F(TreapTest, Init_EmptyTree) {
    EXPECT_TRUE(treap_.Empty());
    EXPECT_EQ(treap_.GetTotalCount(), 0u);
}

// 测试插入单个节点
TEST_F(TreapTest, InsertSingleNode) {
    bool result = treap_.MergeInsert(100, 10, true);
    EXPECT_TRUE(result);

    const Treap::TreapNode* root = treap_.RootNode();
    ASSERT_NE(root, nullptr);
    EXPECT_EQ(root->GetIndex(), 100u);
    EXPECT_EQ(root->GetCount(), 10u);
    EXPECT_EQ(treap_.GetTotalCount(), 10u);
}

// 测试插入相邻节点是否会合并
TEST_F(TreapTest, MergeAdjacentNodes) {
    treap_.MergeInsert(100, 10, true);
    treap_.MergeInsert(110, 5, true); // 应该合并到前一个节点

    const Treap::TreapNode* root = treap_.RootNode();
    ASSERT_NE(root, nullptr);
    EXPECT_EQ(root->GetIndex(), 100u);
    EXPECT_EQ(root->GetCount(), 15u);
    EXPECT_EQ(treap_.GetTotalCount(), 15u);
}

// 测试插入非相邻节点不会合并
TEST_F(TreapTest, NoMergeNonAdjacentNodes) {
    treap_.MergeInsert(100, 10, true);
    treap_.MergeInsert(120, 5, true); // 不相邻，不应合并

    const Treap::TreapNode* root = treap_.RootNode();
    ASSERT_NE(root, nullptr);
    EXPECT_EQ(root->GetIndex(), 100u);
    EXPECT_EQ(root->GetCount(), 10u);

    const Treap::TreapNode* right = root->r;
    ASSERT_NE(right, nullptr);
    EXPECT_EQ(right->GetIndex(), 120u);
    EXPECT_EQ(right->GetCount(), 5u);
    EXPECT_EQ(treap_.GetTotalCount(), 15u);
}

// 测试从树中取出部分单元
TEST_F(TreapTest, TakeUnitsFromNode) {
    treap_.MergeInsert(200, 20, true);

    uint32_t idx = 0;
    bool result = treap_.TakeUnits(5, idx, true);
    EXPECT_TRUE(result);
    EXPECT_EQ(idx, 200u);

    const Treap::TreapNode* root = treap_.RootNode();
    ASSERT_NE(root, nullptr);
    EXPECT_EQ(root->GetIndex(), 205u);
    EXPECT_EQ(root->GetCount(), 15u);
    EXPECT_EQ(treap_.GetTotalCount(), 15u);
}

// 测试无法从空树中取出单元
TEST_F(TreapTest, TakeFromEmptyTreeFails) {
    uint32_t idx = 0;
    bool result = treap_.TakeUnits(5, idx, true);
    EXPECT_FALSE(result);
}

// 测试尝试取出小于当前节点大小的单元
TEST_F(TreapTest, TakeSmallerUnit) {
    treap_.MergeInsert(300, 20, true);

    uint32_t idx = 0;
    bool result = treap_.TakeUnits(5, idx, true);
    EXPECT_TRUE(result);
    EXPECT_EQ(idx, 300u);

    const Treap::TreapNode* root = treap_.RootNode();
    ASSERT_NE(root, nullptr);
    EXPECT_EQ(root->GetIndex(), 305u);
    EXPECT_EQ(root->GetCount(), 15u);
    EXPECT_EQ(treap_.GetTotalCount(), 15u);
}

// 测试插入零个单元应该失败
TEST_F(TreapTest, InsertZeroUnitsFails) {
    bool result = treap_.MergeInsert(400, 0, true);
    EXPECT_FALSE(result);
}

// 测试遍历器能否正确访问所有节点
TEST_F(TreapTest, IteratorTraversesAllNodes) {
    treap_.MergeInsert(500, 10, true);
    treap_.MergeInsert(510, 5, true);

    Treap::Iterator it(treap_);
    const Treap::TreapNode* node = it.Next();
    int count = 0;

    while (node != nullptr) {
        count++;
        node = it.Next();
    }

    EXPECT_EQ(count, 2); // 应该有两个节点
}